package com.langchain4jspringboot.config;

import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.model.openai.OpenAiStreamingChatModel;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.service.MemoryId;
import dev.langchain4j.service.TokenStream;
import dev.langchain4j.service.UserMessage;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

@Configuration
public class AiConfig {

    // 配置记忆对话
    public interface Assistant {
        String chat(String message);
        // 流式响应
        TokenStream stream(String message);
    }

//    配置记忆对话并隔离记忆对话
    public interface AssistantUnique {
        String chat(@MemoryId int memoryId, @UserMessage String userMessage);
        // 流式响应
        TokenStream stream(@MemoryId int memoryId, @UserMessage String userMessage);
    }

    @Bean
    public AssistantUnique assistantUnique(OpenAiChatModel openAiChatModel,
                                           OpenAiStreamingChatModel aiStreamingChatModel) {
        AssistantUnique assistant = AiServices.builder(AssistantUnique.class)
                .chatLanguageModel(openAiChatModel)
                .streamingChatLanguageModel(aiStreamingChatModel)
                .chatMemoryProvider(memoryId ->
                        MessageWindowChatMemory.builder().maxMessages(10)
                                .id(memoryId).build()
                )
                .build();

        return assistant;
    }

    @Bean
    public Assistant assistant(OpenAiChatModel openAiChatModel,
                               OpenAiStreamingChatModel aiStreamingChatModel) {
        ChatMemory chatMemory = MessageWindowChatMemory.withMaxMessages(15);

        Assistant assistant = AiServices.builder(Assistant.class)
                .chatLanguageModel(openAiChatModel)
                .streamingChatLanguageModel(aiStreamingChatModel)
                .chatMemory(chatMemory)
                .build();

        return assistant;
    }
}